package com.fangdd.logtrace.async;

import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.concurrent.*;

/**
 * Created by sulong on 16-9-7.
 */
public class TraceIdScheduledThreadPoolExecutor extends ScheduledThreadPoolExecutor {
    public TraceIdScheduledThreadPoolExecutor(int corePoolSize) {
        super(corePoolSize);
    }

    public TraceIdScheduledThreadPoolExecutor(int corePoolSize, ThreadFactory threadFactory) {
        super(corePoolSize, threadFactory);
    }

    public TraceIdScheduledThreadPoolExecutor(int corePoolSize, RejectedExecutionHandler handler) {
        super(corePoolSize, handler);
    }

    public TraceIdScheduledThreadPoolExecutor(int corePoolSize, ThreadFactory threadFactory, RejectedExecutionHandler handler) {
        super(corePoolSize, threadFactory, handler);
    }


    @Override
    public ScheduledFuture<?> schedule(Runnable command, long delay, TimeUnit unit) {
        return super.schedule(new TraceRunnable(command), delay, unit);
    }

    @Override
    public <V> ScheduledFuture<V> schedule(Callable<V> callable, long delay, TimeUnit unit) {
        return super.schedule(new TraceCallable<V>(callable), delay, unit);
    }

    @Override
    public ScheduledFuture<?> scheduleAtFixedRate(Runnable command, long initialDelay, long period, TimeUnit unit) {
        return super.scheduleAtFixedRate(new TraceRunnable(command), initialDelay, period, unit);
    }

    @Override
    public ScheduledFuture<?> scheduleWithFixedDelay(Runnable command, long initialDelay, long delay, TimeUnit unit) {
        return super.scheduleWithFixedDelay(new TraceRunnable(command), initialDelay, delay, unit);
    }

    @Override
    public <T> Future<T> submit(Callable<T> task) {
        return super.submit(new TraceCallable<T>(task));
    }

    @Override
    public <T> Future<T> submit(Runnable task, T result) {
        return super.submit(new TraceRunnable(task), result);
    }

    @Override
    public Future<?> submit(Runnable task) {
        return super.submit(new TraceRunnable(task));
    }

    @Override
    public <T> List<Future<T>> invokeAll(Collection<? extends Callable<T>> tasks) throws InterruptedException {
        if (tasks == null || tasks.isEmpty()) {
            return super.invokeAll(tasks);
        }

        List<Callable<T>> newTasks = new ArrayList<Callable<T>>(tasks.size());
        for (Callable<T> callable : tasks) {
            newTasks.add(new TraceCallable<T>(callable));
        }

        return super.invokeAll(newTasks);
    }

    @Override
    public <T> List<Future<T>> invokeAll(Collection<? extends Callable<T>> tasks, long timeout, TimeUnit unit) throws InterruptedException {
        if (tasks == null || tasks.isEmpty()) {
            return super.invokeAll(tasks, timeout, unit);
        }

        List<Callable<T>> newTasks = new ArrayList<Callable<T>>(tasks.size());
        for (Callable<T> callable : tasks) {
            newTasks.add(new TraceCallable<T>(callable));
        }

        return super.invokeAll(newTasks, timeout, unit);
    }

    @Override
    public <T> T invokeAny(Collection<? extends Callable<T>> tasks) throws InterruptedException, ExecutionException {
        if (tasks == null || tasks.isEmpty()) {
            return super.invokeAny(tasks);
        }

        List<Callable<T>> newTasks = new ArrayList<Callable<T>>(tasks.size());
        for (Callable<T> callable : tasks) {
            newTasks.add(new TraceCallable<T>(callable));
        }

        return super.invokeAny(newTasks);
    }

    @Override
    public <T> T invokeAny(Collection<? extends Callable<T>> tasks, long timeout, TimeUnit unit) throws InterruptedException, ExecutionException, TimeoutException {
        if (tasks == null || tasks.isEmpty()) {
            return super.invokeAny(tasks, timeout, unit);
        }

        List<Callable<T>> newTasks = new ArrayList<Callable<T>>(tasks.size());
        for (Callable<T> callable : tasks) {
            newTasks.add(new TraceCallable<T>(callable));
        }

        return super.invokeAny(newTasks, timeout, unit);
    }

    @Override
    public void execute(Runnable command) {
        super.execute(new TraceRunnable(command));
    }
}
